from transformers import AutoTokenizer, AutoModelForSequenceClassification
from langchain.tools import Tool
from typing import Optional
from langchain.agents import initialize_agent, AgentType
from langchain.llms import OpenAI  # 或其他LangChain支持的LLM

# 加载FinBERT情感分析模型
model_name = "yiyanghkust/finbert-tone"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def finbert_sentiment_analysis(text: str) -> str:
    """使用FinBERT分析金融文本情感"""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model(**inputs)
    predictions = outputs.logits.argmax(dim=-1)
    labels = ["Positive", "Negative", "Neutral"]
    return labels[predictions.item()]

# 注册为Tool
finbert_tool = Tool(
    name="FinBERT_Sentiment_Analysis",
    func=finbert_sentiment_analysis,
    description="Useful for analyzing sentiment of financial news or reports. Input a text string, output 'Positive', 'Negative', or 'Neutral'."
)

# 选择基础LLM（如GPT-3.5）
llm = OpenAI(temperature=0)  

# 定义工具列表
tools = [finbert_tool]  # 可加入其他工具如Search、Calculator

# 初始化Agent
agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # 适用于简单任务
    verbose=True
)

query = "Analyze the sentiment of this text: 'The company's revenue growth exceeded expectations, but costs rose sharply.'"
result = agent.run(query)
print(result)